﻿#pragma once
#include "Types.h"
#include <vtkm/cont/ExecutionAndControlObjectBase.h>
//#include <vtkm/cont/ExecutionObjectBase.h>

struct EquationSet : vtkm::cont::ExecutionAndControlObjectBase
{
  static constexpr UInt8 NUM_VARIABLES = 5;
  using ValueType = Vec<Real, NUM_VARIABLES>;

  VTKM_CONT EquationSet(Real mach)
    : _mach(mach)
  {
  }

  VTKM_EXEC_CONT ValueType PrimaryToConservation(const ValueType& primary_value) const
  {
    ValueType value;
    auto r = primary_value[0];
    auto u = primary_value[1];
    auto v = primary_value[2];
    auto w = primary_value[3];
    auto p = primary_value[4];

    value[0] = r;
    value[1] = r * u;
    value[2] = r * v;
    value[3] = r * w;
    value[4] = p / _gamm1 + 0.5 * r * (u * u + v * v);
    return value;
  }

  template<typename CoordType>
  VTKM_EXEC_CONT ValueType InitialCondition(const CoordType& p)
  {
    ValueType primary_value;
    //primary_value = ValueType{ 1.4, 0.3, 0, 0, 1 };
    //Real pre = 1 / _gamma / _mach / _mach;
    primary_value = ValueType{ 1, 1, 0, 0, 1 / _gamma / _mach / _mach };

    return PrimaryToConservation(primary_value);
  }


  template<typename CoordType>
  VTKM_EXEC_CONT ValueType
  BoundaryCondition(const UInt8& bid, const ValueType& ul, const Vec3f& normal, const CoordType& p)
  {
    ValueType ur;
    ur = ul;
    if (bid == 1) // 出口
    {
      ur[0] = ul[0];
      ur[1] = ul[1];
      ur[2] = ul[2];
      ur[3] = ul[3];
      ur[4] = ul[4];
    }
    else if (bid == 2) // 进口
    {
      ValueType primary_value;
      primary_value = ValueType{ 1, 1, 0, 0, 1 / _gamma / _mach / _mach };

      ValueType value;
      auto r = primary_value[0];
      auto u = primary_value[1];
      auto v = primary_value[2];
      auto w = primary_value[3];
      auto pre = primary_value[4];

      ur[0] = r;
      ur[1] = r * u;
      ur[2] = r * v;
      ur[3] = r * w;
      ur[4] = pre / _gamm1 + 0.5 * r * (u * u + v * v);
    }
    else if (bid == 3) //wall
    {
      Vec3f mom{ ul[1], ul[2], ul[3] };
      Real vn = vtkm::Dot(mom, normal);
      ur[0] = ul[0];
      ur[1] = ul[1] - 2 * vn * normal[0];
      ur[2] = ul[2] - 2 * vn * normal[1];
      ur[3] = ul[3] - 2 * vn * normal[2];
      ur[4] = ul[4];
    }
    else
    {
      // TODO: Raise Error.
      //this->RaiseError("不可用的BC Type.");
    }

    return ur;
  }

  VTKM_EXEC_CONT
  void InvisTerm(const ValueType& uh, Vec3f* inv_term)
  {
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };
    FloatDefault p = _gamm1 * (uh[4] - 0.5 * uh[0] * vtkm::Dot(vel, vel));
    FloatDefault h = (uh[4] + p) / uh[0];

    inv_term[0] = uh[0] * vel;

    inv_term[1] = uh[1] * vel;
    inv_term[1][0] += p;

    inv_term[2] = uh[2] * vel;
    inv_term[2][1] += p;

    inv_term[3] = uh[3] * vel;
    inv_term[3][2] += p;

    inv_term[4] = uh[0] * h * vel;
  }

  VTKM_EXEC_CONT
  void InvisTermFlux(const ValueType& uh, const Vec3f& normal, ValueType& flux)
  {
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };
    FloatDefault p = _gamm1 * (uh[4] - 0.5 * uh[0] * vtkm::Dot(vel, vel));
    FloatDefault h = (uh[4] + p) / uh[0];

    Real vn = vtkm::Dot(vel, normal);
    flux[0] = uh[0] * vn;
    flux[1] = uh[1] * vn + p * normal[0];
    flux[2] = uh[2] * vn + p * normal[1];
    flux[3] = uh[3] * vn + p * normal[2];
    flux[4] = uh[0] * h * vn;
  }

  VTKM_EXEC_CONT void RiemannFlux(const ValueType& ul,
                                  const ValueType& ur,
                                  const Vec3f& normal,
                                  ValueType& riemann_flux)
  {
    ValueType flux_l, flux_r;
    InvisTermFlux(ul, normal, flux_l);
    InvisTermFlux(ur, normal, flux_r);
    riemann_flux = (flux_l + flux_r) / 2. + 2 * (ul - ur);
  }

  EquationSet PrepareForExecution(vtkm::cont::DeviceAdapterId, vtkm::cont::Token&) const
  {
    return *this;
  }

  EquationSet PrepareForControl() const { return *this; }

  Real _mach;
  Real _gamma = 1.4, _gamm1 = _gamma - 1;
};
